# coding = utf-8
import os
import cv2
from PIL import Image
import numpy as np


def getImagesAndLabels(path,detector):
    imagePaths = [os.path.join(path, f) for f in os.listdir(path)]
    faceSamples = []
    ids = []
    if not imagePaths:
        return -1,-1
    for imagePath in imagePaths:
        PIL_img = Image.open(imagePath).convert('L')
        img_numpy = np.array(PIL_img, 'uint8')
        id = int(os.path.split(imagePath)[-1].split(".")[1])
        print(imagePath,id)
        faces = detector.detectMultiScale(img_numpy)
        for (x, y, w, h) in faces:
            faceSamples.append(img_numpy[y:y + h, x: x + w])
            ids.append(id)
    return faceSamples, ids

def trainModel():
    path = './FaceData/'
    recognizer = cv2.face.LBPHFaceRecognizer_create()

    detector = cv2.CascadeClassifier(r'./CV2Data/haarcascade_frontalface_default.xml')
    print('Training faces. It will take a few seconds. Waiting...')
    faces, ids = getImagesAndLabels(path,detector)
    if faces == -1 or ids == -1:
        return -1
    recognizer.train(faces, np.array(ids))
    print('Training has finished!')
    recognizer.write(r'./Model/trainer-2021.yml')
    print("{0} faces trained. Exiting Program.".format(len(np.unique(ids))))
    return 0

